import os
import numpy as np
import cv2
import onnxruntime
from utils.utils import image_transform

directory_name = './resources/face_datas'
model_path = './resources/model.onnx'
save_images_path = './resources/images_path.txt'
end_path = './resources/face_feature_vector.txt'


def read_directory(directory_name):
    images_path = []
    array_of_img = []
    for filename in os.listdir(r"./" + directory_name):
        img = cv2.imread(directory_name + "/" + filename)
        images_path.append(directory_name + "/" + filename)
        img = image_transform(img)
        img_dict = {'input': img}
        array_of_img.append(img_dict)


def create_txt(imgs_list):
    face_feature_vector = []
    ort_session = onnxruntime.InferenceSession(model_path)

    for img_dict in imgs_list:
        ort_output = ort_session.run(['output'], img_dict)[0]
        face_feature_vector.append(ort_output.ravel())

    face_feature_vector = np.array(face_feature_vector)
    # 将人脸图片库的facenet运算后的128维特征向量保存到face_feature_vector.txt
    np.savetxt(end_path, face_feature_vector, )


def write_images_path(path):
    f = open(save_images_path, 'w')
    for line in path:
        f.write(line + '\n')
    f.close()


if __name__ == '__main__':
    read_directory(directory_name)
    write_images_path(images_path)
    create_txt(array_of_img)
